#!/usr/bin/python

import logging
import random


class AssignmentPriorityQueue(object):
  """Stateful priority queue that handles scheduling with constraints."""

  def __init__(self, size, max_counts, max_night_counts, schedule):
    self.count_size = size
    self.counts = [0] * size  # Number of shifts per doctor
    self.night_counts = [0] * size  # Number of shifts per doctor
    self.max_counts = max_counts  # Each doctor's maximum
    self.max_night_counts = max_night_counts
    self.schedule = schedule
    self._last_shift_asst = []  # State from last assignment
    self._last_day_shift_asst = [0] * size

  def _ChooseNextCandidates(self, candidates, counts, max_counts, 
      consecutive_counts, excludes):
    smallest = 1000000  # Sentinel
    new_candidates = []

    # Filter all valid candidates.
    for index in candidates:
      if (counts[index] < max_counts[index] and
          index not in excludes):
        new_candidates.append(index)
        smallest = min(smallest, counts[index])
    if not new_candidates:
      logging.info("No candidates after first stage filter.")
      return None

    # Try to place consecutive days.
    filtered_consecutive = False
    if consecutive_counts:
      max_consecutive_so_far = max(consecutive_counts)
      if max_consecutive_so_far:
        filtered_by_consecutive_days = new_candidates[:]
        for index in new_candidates:
          threshhold = random.randint(3, 5)
          consecutive_days = consecutive_counts[index]
          if consecutive_days == 0 or consecutive_days >= threshhold:
            filtered_by_consecutive_days.remove(index)
        if filtered_by_consecutive_days:
          new_candidates = filtered_by_consecutive_days
          filtered_consecutive = True
    if not new_candidates:
      logging.info("No candidates after consecutive day filter.")
      return None

    if not filtered_consecutive:
      # Choose the candidate among the valid with the smallest
      # number of shifts already assigned. This heuristic helps
      # to meet minimum shift requirements.
      new_candidates = [index for index in new_candidates
          if counts[index] == smallest]

    return new_candidates

  def NextAvailable(self, shift, count, excludes):

    def Concat(*lists):
      l = []
      for e in lists:
        l.extend(e)
      return l

    chosen = []

    for i in xrange(count):
      # TODO Inclusion list would replace range here.
      candidates = range(self.count_size)

      if self.schedule.IsNight(shift):
        all_excludes = Concat(self._last_shift_asst, chosen, excludes)
        candidates = self._ChooseNextCandidates(candidates, self.night_counts,
            self.max_night_counts, None, all_excludes)
      if not candidates:
        return None

      if self.schedule.IsDay(shift):
        consecutive_counts = self._last_day_shift_asst
      else:
        consecutive_counts = None
      all_excludes = Concat(self._last_shift_asst, chosen, excludes)
      candidates = self._ChooseNextCandidates(candidates, self.counts,
          self.max_counts, consecutive_counts, all_excludes)
      if not candidates:
        return None

      next_index = random.choice(candidates)
      self.counts[next_index] += 1
      if self.schedule.IsNight(shift):
        self.night_counts[next_index] += 1
      chosen.append(next_index)

    self._last_shift_asst = chosen[:]

    if self.schedule.IsDay(shift):
      for i in range(self.count_size):
        if i in chosen:
          self._last_day_shift_asst[next_index] += 1
        else:
          self._last_day_shift_asst[i] = 0

    return chosen
